//SPDX-FileCopyrightText: Copyright 2022-2024 深圳市同心圆网络有限公司
//SPDX-License-Identifier: GPL-3.0-only

package impl

import (
	"context"
	"encoding/base64"
	"io"
	"sync/atomic"
	"time"

	"atomgit.com/openlinksaas-org/ssh-proxy/gen_code"
	"github.com/google/uuid"
	"golang.org/x/crypto/ssh"
)

type ShellInfo struct {
	session  *ssh.Session
	Index    int32
	recvChan chan []byte
	stdin    io.WriteCloser
}

type ShellApiImpl struct {
}

func (s *ShellApiImpl) ShellChangePost(ctx context.Context, req gen_code.ShellChangePostRequest) (gen_code.ImplResponse, error) {
	if req.Token != token {
		return gen_code.ImplResponse{
			Code: 500,
			Body: &gen_code.ErrInfo{
				ErrMsg: "wrong token",
			},
		}, nil
	}
	info, ok := clientMap.Load(req.Id)
	if !ok {
		return gen_code.ImplResponse{
			Code: 500,
			Body: &gen_code.ErrInfo{
				ErrMsg: "miss client",
			},
		}, nil
	}
	client := info.(*ClientInfo)
	shellInfo, ok := client.shellMap.Load(req.ShellId)
	if !ok {
		return gen_code.ImplResponse{
			Code: 500,
			Body: &gen_code.ErrInfo{
				ErrMsg: "miss shell",
			},
		}, nil
	}
	sess := shellInfo.(*ShellInfo).session
	err := sess.WindowChange(int(req.Height), int(req.Width))
	if err != nil {
		return gen_code.ImplResponse{
			Code: 500,
			Body: &gen_code.ErrInfo{
				ErrMsg: err.Error(),
			},
		}, nil
	}
	return gen_code.ImplResponse{
		Code: 200,
		Body: req.ShellId,
	}, nil
}

func (s *ShellApiImpl) ShellClosePost(ctx context.Context, req gen_code.ShellClosePostRequest) (gen_code.ImplResponse, error) {
	if req.Token != token {
		return gen_code.ImplResponse{
			Code: 500,
			Body: &gen_code.ErrInfo{
				ErrMsg: "wrong token",
			},
		}, nil
	}
	info, ok := clientMap.Load(req.Id)
	if !ok {
		return gen_code.ImplResponse{
			Code: 500,
			Body: &gen_code.ErrInfo{
				ErrMsg: "miss client",
			},
		}, nil
	}
	client := info.(*ClientInfo)
	shellInfo, ok := client.shellMap.Load(req.ShellId)
	if !ok {
		return gen_code.ImplResponse{
			Code: 500,
			Body: &gen_code.ErrInfo{
				ErrMsg: "miss shell",
			},
		}, nil
	}
	sess := shellInfo.(*ShellInfo).session
	sess.Close()
	close(shellInfo.(*ShellInfo).recvChan)
	client.shellMap.Delete(req.ShellId)
	return gen_code.ImplResponse{
		Code: 200,
		Body: req.ShellId,
	}, nil
}

func (s *ShellApiImpl) ShellListPost(ctx context.Context, req gen_code.PortforwardListPostRequest) (gen_code.ImplResponse, error) {
	if req.Token != token {
		return gen_code.ImplResponse{
			Code: 500,
			Body: &gen_code.ErrInfo{
				ErrMsg: "wrong token",
			},
		}, nil
	}
	info, ok := clientMap.Load(req.Id)
	if !ok {
		return gen_code.ImplResponse{
			Code: 500,
			Body: &gen_code.ErrInfo{
				ErrMsg: "miss client",
			},
		}, nil
	}
	client := info.(*ClientInfo)
	retList := []gen_code.ShellInfo{}
	client.shellMap.Range(func(key, value any) bool {
		retList = append(retList, gen_code.ShellInfo{
			ShellId:    key.(string),
			ShellIndex: value.(*ShellInfo).Index,
		})
		return true
	})
	return gen_code.ImplResponse{
		Code: 200,
		Body: retList,
	}, nil
}

func (s *ShellApiImpl) ShellOpenPost(ctx context.Context, req gen_code.ShellOpenPostRequest) (gen_code.ImplResponse, error) {
	if req.Token != token {
		return gen_code.ImplResponse{
			Code: 500,
			Body: &gen_code.ErrInfo{
				ErrMsg: "wrong token",
			},
		}, nil
	}
	info, ok := clientMap.Load(req.Id)
	if !ok {
		return gen_code.ImplResponse{
			Code: 500,
			Body: &gen_code.ErrInfo{
				ErrMsg: "miss client",
			},
		}, nil
	}
	client := info.(*ClientInfo)
	sess, err := client.client.realClient.NewSession()
	if err != nil {
		return gen_code.ImplResponse{
			Code: 500,
			Body: &gen_code.ErrInfo{
				ErrMsg: err.Error(),
			},
		}, nil
	}

	err = sess.RequestPty(req.TermType, int(req.Height), int(req.Width), ssh.TerminalModes{
		ssh.ECHO:          1,
		ssh.TTY_OP_ISPEED: 14400,
		ssh.TTY_OP_OSPEED: 14400,
	})
	if err != nil {
		sess.Close()
		return gen_code.ImplResponse{
			Code: 500,
			Body: &gen_code.ErrInfo{
				ErrMsg: err.Error(),
			},
		}, nil
	}

	stdout, err := sess.StdoutPipe()
	if err != nil {
		sess.Close()
		return gen_code.ImplResponse{
			Code: 500,
			Body: &gen_code.ErrInfo{
				ErrMsg: err.Error(),
			},
		}, nil
	}

	stderr, err := sess.StderrPipe()
	if err != nil {
		sess.Close()
		return gen_code.ImplResponse{
			Code: 500,
			Body: &gen_code.ErrInfo{
				ErrMsg: err.Error(),
			},
		}, nil
	}

	stdin, err := sess.StdinPipe()
	if err != nil {
		sess.Close()
		return gen_code.ImplResponse{
			Code: 500,
			Body: &gen_code.ErrInfo{
				ErrMsg: err.Error(),
			},
		}, nil
	}

	err = sess.Shell()
	if err != nil {
		sess.Close()
		return gen_code.ImplResponse{
			Code: 500,
			Body: &gen_code.ErrInfo{
				ErrMsg: err.Error(),
			},
		}, nil
	}

	ch := make(chan []byte, 1000)

	go s.readShellData(ch, stdout)
	go s.readShellData(ch, stderr)

	shellId := uuid.NewString()
	shellIndex := atomic.AddInt32(&client.shellIndex, 1)
	client.shellMap.Store(shellId, &ShellInfo{
		session:  sess,
		Index:    shellIndex,
		recvChan: ch,
		stdin:    stdin,
	})

	return gen_code.ImplResponse{
		Code: 200,
		Body: &gen_code.ShellInfo{
			ShellId:    shellId,
			ShellIndex: shellIndex,
		},
	}, nil
}

func (s *ShellApiImpl) readShellData(ch chan []byte, reader io.Reader) {
	defer func() {
		recover()
	}()

	for {
		buf := make([]byte, 4*1024)
		n, err := reader.Read(buf)
		if n > 0 {
			ch <- buf[:n]
		}
		if err != nil {
			break
		}
	}
}

func (s *ShellApiImpl) ShellReadPost(ctx context.Context, req gen_code.ShellClosePostRequest) (gen_code.ImplResponse, error) {
	defer func() {
		recover()
	}()

	if req.Token != token {
		return gen_code.ImplResponse{
			Code: 500,
			Body: &gen_code.ErrInfo{
				ErrMsg: "wrong token",
			},
		}, nil
	}
	info, ok := clientMap.Load(req.Id)
	if !ok {
		return gen_code.ImplResponse{
			Code: 500,
			Body: &gen_code.ErrInfo{
				ErrMsg: "miss client",
			},
		}, nil
	}
	client := info.(*ClientInfo)
	shellInfo, ok := client.shellMap.Load(req.ShellId)
	if !ok {
		return gen_code.ImplResponse{
			Code: 500,
			Body: &gen_code.ErrInfo{
				ErrMsg: "miss shell",
			},
		}, nil
	}
	select {
	case data := <-shellInfo.(*ShellInfo).recvChan:
		return gen_code.ImplResponse{
			Code: 200,
			Body: base64.StdEncoding.EncodeToString(data),
		}, nil
	case <-time.After(10 * time.Second):
		return gen_code.ImplResponse{
			Code: 200,
			Body: "",
		}, nil
	}
}

func (s *ShellApiImpl) ShellWritePost(ctx context.Context, req gen_code.ShellWritePostRequest) (gen_code.ImplResponse, error) {
	if req.Token != token {
		return gen_code.ImplResponse{
			Code: 500,
			Body: &gen_code.ErrInfo{
				ErrMsg: "wrong token",
			},
		}, nil
	}
	info, ok := clientMap.Load(req.Id)
	if !ok {
		return gen_code.ImplResponse{
			Code: 500,
			Body: &gen_code.ErrInfo{
				ErrMsg: "miss client",
			},
		}, nil
	}
	client := info.(*ClientInfo)
	shellInfo, ok := client.shellMap.Load(req.ShellId)
	if !ok {
		return gen_code.ImplResponse{
			Code: 500,
			Body: &gen_code.ErrInfo{
				ErrMsg: "miss shell",
			},
		}, nil
	}
	data, err := base64.StdEncoding.DecodeString(req.Data)
	if err != nil {
		return gen_code.ImplResponse{
			Code: 500,
			Body: &gen_code.ErrInfo{
				ErrMsg: err.Error(),
			},
		}, nil
	}

	_, err = shellInfo.(*ShellInfo).stdin.Write(data)
	if err != nil {
		return gen_code.ImplResponse{
			Code: 500,
			Body: &gen_code.ErrInfo{
				ErrMsg: err.Error(),
			},
		}, nil
	}
	return gen_code.ImplResponse{
		Code: 200,
		Body: req.ShellId,
	}, nil
}
